{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pickle\n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "import random\n",
    "    \n",
    "# load source words\n",
    "source_words_path = os.path.join(os.getcwd(), 'source_words.pkl')\n",
    "with open(source_words_path, 'rb') as f_source_words:\n",
    "    source_words = pickle.load(f_source_words)\n",
    "    \n",
    "# load target words\n",
    "target_words_path = os.path.join(os.getcwd(), 'target_words.pkl')\n",
    "with open(target_words_path, 'rb') as f_target_words:\n",
    "    target_words = pickle.load(f_target_words)\n",
    "    \n",
    "# load label words\n",
    "label_words_path = os.path.join(os.getcwd(), 'label_words.pkl')\n",
    "with open(label_words_path, 'rb') as f_label_words:\n",
    "    label_words = pickle.load(f_label_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "945\n",
      "133\n",
      "27\n",
      "1\n",
      "3\n",
      "2\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "print(len(source_words))\n",
    "print(len(target_words))\n",
    "print(len(label_words))\n",
    "print(source_words['<pad>'])\n",
    "print(source_words['<eos>'])\n",
    "print(source_words['<sos>'])\n",
    "print(source_words['<unk>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 构建attention权重计算方式\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, hidden_dim):\n",
    "        super(Attention, self).__init__()\n",
    "        self.attn = nn.Linear((hidden_dim * 2), hidden_dim)\n",
    "        self.v = nn.Linear(hidden_dim, 1, bias=False)\n",
    "\n",
    "    def concat_score(self, hidden, encoder_output):\n",
    "        seq_len = encoder_output.shape[1]\n",
    "        hidden = hidden.unsqueeze(1).repeat(1, seq_len, 1) # [batch_size, seq_len, hidden_size]\n",
    "        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_output),dim=2))) # [batch_size, seq_len, hidden_dim]\n",
    "        attention = self.v(energy).squeeze(2) #[batch_size, seq_len]\n",
    "        return attention #[batch_size, seq_len]\n",
    "\n",
    "    def forward(self, hidden, encoder_output):\n",
    "        # hidden = [batch_size, hidden_size]\n",
    "        # #encoder_output=[batch_size, seq_len, hidden_size]\n",
    "        \n",
    "        attn_energies = self.concat_score(hidden, encoder_output)\n",
    "\n",
    "        return F.softmax(attn_energies, dim=1).unsqueeze(1) #softmax归一化，[batch_size, 1, seq_len]\n",
    "    \n",
    "#构建模型\n",
    "class BirnnAttention(nn.Module):\n",
    "    def __init__(self, source_input_dim, source_emb_dim, hidden_dim, n_layers, dropout, pad_index, slot_output_size, intent_output_size, slot_embed_dim, predict_flag):\n",
    "        super(BirnnAttention, self).__init__()\n",
    "        self.pad_index = pad_index\n",
    "        self.hidden_dim = hidden_dim//2 # 双向lstm\n",
    "        self.n_layers = n_layers\n",
    "        self.slot_output_size = slot_output_size\n",
    "        # 是否预测模式\n",
    "        self.predict_flag = predict_flag\n",
    "        \n",
    "        self.source_embedding = nn.Embedding(source_input_dim, source_emb_dim, padding_idx=pad_index)\n",
    "        # 双向gru，隐层维度是hidden_dim\n",
    "        self.source_gru = nn.GRU(source_emb_dim, self.hidden_dim, n_layers, dropout=dropout, bidirectional=True, batch_first=True) #使用双向\n",
    "        \n",
    "        \n",
    "        # 单个cell的隐层维度与gru隐层维度一样，为hidden_dim\n",
    "        self.gru_cell = nn.GRUCell(slot_embed_dim + (2 * hidden_dim), hidden_dim)\n",
    "        self.attention = Attention(hidden_dim)\n",
    "        # 意图intent预测\n",
    "        self.intent_output = nn.Linear(hidden_dim * 2, intent_output_size)\n",
    "        # 槽slot预测\n",
    "        self.slot_output = nn.Linear(hidden_dim, slot_output_size)\n",
    "        self.slot_embedding = nn.Embedding(slot_output_size, slot_embed_dim)\n",
    "        \n",
    "    def forward(self, source_input, source_len):\n",
    "        '''\n",
    "        source_input:[batch_size, seq_len]\n",
    "        source_len:[batch_size]\n",
    "        '''\n",
    "        if self.predict_flag:\n",
    "            assert len(source_input) == 1, '预测时一次输入一句话'\n",
    "            seq_len = source_len[0]\n",
    "            \n",
    "            # 1.Encoder阶段，将输入的source进行编码\n",
    "            # source_embedded:[batch_size, seq_len, source_emb_dim]\n",
    "            source_embedded = self.source_embedding(source_input)\n",
    "            packed = torch.nn.utils.rnn.pack_padded_sequence(source_embedded, source_len, batch_first=True, enforce_sorted=True) #这里enfore_sotred=True要求数据根据词数排序\n",
    "            source_output, hidden = self.source_gru(packed)\n",
    "            # source_output=[batch_size, seq_len, 2 * self.hidden_size]，这里的2*self.hidden_size = hidden_dim\n",
    "            # hidden=[n_layers * 2, batch_size, self.hidden_size]\n",
    "            source_output, _ = torch.nn.utils.rnn.pad_packed_sequence(source_output, batch_first=True, padding_value=self.pad_index, total_length=len(source_input[0])) #这个会返回output以及压缩后的legnths\n",
    "            '''\n",
    "            source_hidden[-2,:,:]是gru最后一步的forward\n",
    "            source_hidden[-1,:,:]是gru最后一步的backward\n",
    "            '''\n",
    "            # source_hidden=[batch_size, 2*self.hidden_size]\n",
    "            source_hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)\n",
    "            #保存注意力向量\n",
    "            attention_context = torch.zeros(1, seq_len, self.hidden_dim * 2)\n",
    "            output_tokens = []\n",
    "           \n",
    "            aligns = source_output.transpose(0,1) #对齐向量\n",
    "            \n",
    "            input = torch.tensor(2).unsqueeze(0)  # 预测阶段解码器输入第一个token-> <sos>\n",
    "            for s in range(1, seq_len):\n",
    "                aligned = aligns[s].unsqueeze(1)# [batch_size, 1, hidden_size*2]\n",
    "                 # embedded=[1, 1, slot_embed_dim]\n",
    "                slot_embedded = self.slot_embedding(input)\n",
    "                slot_embedded = slot_embedded.unsqueeze(0)\n",
    "                # 利用利用上一步的hidden与encoder_output，计算attention权重\n",
    "                # attention_weights=[batch_size, 1, seq_len]\n",
    "                attention_weights = self.attention(source_hidden, source_output)\n",
    "\n",
    "                '''\n",
    "                以下是计算上下文：利用attention权重与encoder_output计算attention上下文向量\n",
    "                注意力权重分布用于产生编码器隐藏状态的加权和，加权平均的过程。得到的向量称为上下文向量\n",
    "                '''\n",
    "                context = attention_weights.bmm(source_output)\n",
    "                attention_context[:,s,:] = context\n",
    "              \n",
    "                combined_grucell_input = torch.cat([aligned, slot_embedded, context], dim =2)\n",
    "             \n",
    "                source_hidden = self.gru_cell(combined_grucell_input.squeeze(1), source_hidden)\n",
    "              \n",
    "                slot_prediction = self.slot_output(source_hidden)\n",
    "            \n",
    "                input = slot_prediction.argmax(1)\n",
    "                output_token = input.squeeze().detach().item()\n",
    "               \n",
    "                output_tokens.append(output_token)\n",
    "            \n",
    "             #意图识别\n",
    "            #拼接注意力向量和encoder的输出\n",
    "            combined_attention_sourceoutput = torch.cat([attention_context, source_output], dim=2)\n",
    "            intent_outputs = self.intent_output(torch.mean(combined_attention_sourceoutput, dim = 1))\n",
    "            intent_outputs = intent_outputs.squeeze()\n",
    "            intent_outputs = intent_outputs.argmax()\n",
    "            return output_tokens, intent_outputs\n",
    "        \n",
    "        else:\n",
    "            # 1.Encoder阶段，将输入的source进行编码\n",
    "            # source_embedded:[batch_size, seq_len, source_emb_dim]\n",
    "            source_embedded = self.source_embedding(source_input)\n",
    "            packed = torch.nn.utils.rnn.pack_padded_sequence(source_embedded, source_len, batch_first=True, enforce_sorted=True) #这里enfore_sotred=True要求数据根据词数排序\n",
    "            source_output, hidden = self.source_gru(packed)\n",
    "            # source_output=[batch_size, seq_len, 2 * self.hidden_size]，这里的2*self.hidden_size = hidden_dim\n",
    "            # hidden=[n_layers * 2, batch_size, self.hidden_size]\n",
    "            source_output, _ = torch.nn.utils.rnn.pad_packed_sequence(source_output, batch_first=True, padding_value=self.pad_index, total_length=len(source_input[0])) #这个会返回output以及压缩后的legnths\n",
    "            '''\n",
    "            source_hidden[-2,:,:]是gru最后一步的forward\n",
    "            source_hidden[-1,:,:]是gru最后一步的backward\n",
    "            '''\n",
    "            # source_hidden=[batch_size, 2*self.hidden_size]\n",
    "            source_hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)\n",
    "\n",
    "\n",
    "            # 2.Decoder阶段，预测slot与intent\n",
    "            batch_size = source_input.shape[0]\n",
    "            seq_len = source_input.shape[1]\n",
    "            # 保存slot的预测概率\n",
    "            slot_outputs = torch.zeros(batch_size, seq_len, self.slot_output_size).to(device)\n",
    "\n",
    "            #保存注意力向量\n",
    "            attention_context = torch.zeros(batch_size, seq_len, self.hidden_dim * 2).to(device)\n",
    "\n",
    "            # 每个batch数据的第一个字符<sos>对应的是index是2\n",
    "            input = torch.tensor(2).repeat(batch_size).to(device)\n",
    "            aligns = source_output.transpose(0,1) # 利用encoder output最后一层的每一个时间步\n",
    "            # 槽识别\n",
    "            for t in range(1, seq_len):\n",
    "                '''\n",
    "                解码器输入的初始hidden为encoder的最后一步的hidden\n",
    "                接收输出即predictions和新的hidden状态\n",
    "                '''\n",
    "                aligned = aligns[t].unsqueeze(1)# [batch_size, 1, hidden_size] # hidden_size包含前向和后向隐状态向量\n",
    "                input = input.unsqueeze(1)\n",
    "                # input=[batch_size, 1]\n",
    "                # hidden=[batch_size, hidden_size] 初始化为encoder的最后一层 [batch_size, hidden_size]\n",
    "                # encoder_output=[batch_size, seq_len, hidden_dim*2]\n",
    "                # aligned=[batch_size, 1, hidden_dim*2]\n",
    "\n",
    "                # embedded=[batch_sze, 1, slot_embed_dim]\n",
    "                slot_embedded = self.slot_embedding(input)\n",
    "\n",
    "                # 利用利用上一步的hidden与encoder_output，计算attention权重\n",
    "                # attention_weights=[batch_size, 1, seq_len]\n",
    "                attention_weights = self.attention(source_hidden, source_output)\n",
    "\n",
    "                '''\n",
    "                以下是计算上下文：利用attention权重与encoder_output计算attention上下文向量\n",
    "                注意力权重分布用于产生编码器隐藏状态的加权和，加权平均的过程。得到的向量称为上下文向量\n",
    "                '''\n",
    "                context = attention_weights.bmm(source_output) # [batch_size, 1, seq_len] * [batch_size, seq_len, hidden_dim]=[batch_size, 1, hidden_dim]\n",
    "                attention_context[:,t,:] = context.squeeze(1)\n",
    "                #combined_grucell_input=[batch_size, 1, (hidden_size + slot_embed_dim + hidden_dim)]\n",
    "                combined_grucell_input = torch.cat([aligned, slot_embedded, context], dim =2)\n",
    "                # [batch_size, hidden_dim]\n",
    "                source_hidden = self.gru_cell(combined_grucell_input.squeeze(1), source_hidden)\n",
    "                # 预测slot, [batch_size, slot_output_size]\n",
    "                slot_prediction = self.slot_output(source_hidden)\n",
    "                slot_outputs[:, t, :] = slot_prediction\n",
    "                # 获取预测的最大概率的token\n",
    "                input = slot_prediction.argmax(1)\n",
    "            #意图识别\n",
    "            #拼接注意力向量和encoder的输出，[batch_size, seq_len, hidden_dim * 2]\n",
    "            combined_attention_sourceoutput = torch.cat([attention_context, source_output], dim=2)\n",
    "            intent_outputs = self.intent_output(torch.mean(combined_attention_sourceoutput, dim = 1))\n",
    "\n",
    "            return slot_outputs, intent_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\rnn.py:50: UserWarning: dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.5 and num_layers=1\n",
      "  \"num_layers={}\".format(dropout, num_layers))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['which', 'airlines', 'have', 'first', 'class', 'flights', 'today', '<eos>']\n",
      "[58, 31, 73, 46, 61, 6, 281, 3]\n",
      "slot_prediciton:o o b-class_type i-class_type o b-depart_date.today_relative <eos>\n",
      "intent_prediction:airline\n"
     ]
    }
   ],
   "source": [
    "source_emb_dim = 64\n",
    "slot_embed_dim = 64\n",
    "hidden_dim = 128\n",
    "n_layers = 1\n",
    "dropout = 0.5\n",
    "\n",
    "model_path = os.path.join(os.getcwd(), \"model.h5\")\n",
    "\n",
    "input_dim = len(source_words) # source 词典大小（即词数量）\n",
    "output_dim = len(target_words) # target 词典大小（即实体类型数量）\n",
    "label_dim = len(label_words) # label 词典大小（即意图类别数量）\n",
    "\n",
    "model = BirnnAttention(input_dim, source_emb_dim, hidden_dim, n_layers, dropout, source_words['<pad>'], output_dim, label_dim, slot_embed_dim, True)\n",
    "   \n",
    "model.load_state_dict(torch.load(model_path))\n",
    "model.eval()\n",
    "\n",
    "sentence = \"i would like to find a flight from charlotte to las vegas that makes a stop in st. louis\"\n",
    "sentence2 = \"which airlines have first class flights today\"\n",
    "with torch.no_grad():\n",
    "    tokenized = sentence2.split()  # tokenize the sentence\n",
    "    tokenized.append('<eos>')\n",
    "    indexed = [source_words[t] for t in tokenized]  # convert to integer sequence\n",
    "    print(tokenized)\n",
    "    print(indexed)\n",
    "    tensor = torch.LongTensor(indexed)  # convert to tensor\n",
    "    tensor = tensor.unsqueeze(0)  # reshape in form of batch,no. of words\n",
    "    slot_outputs, intent_outputs = model(tensor, [len(tensor[0])])  # prediction\n",
    "    intent = intent_outputs.detach().item()\n",
    "    slot_prediction = [target_words.itos[t] for t in slot_outputs]\n",
    "\n",
    "    print('slot_prediciton:{}'.format(' '.join(slot_prediction)))\n",
    "    print('intent_prediction:{}'.format(label_words.itos[intent]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
